#############################################################################################
#############################################################################################
# 
#   Open Source License/Disclaimer, Forecast Systems Laboratory NOAA/OAR/GSD, 
#   325 Broadway Boulder, CO 80305
#
#   This software is distributed under the Open Source Definition, which may be 
#   found at http://www.opensource.org/.
#   In particular, redistribution and use in source and binary forms, with or 
#   without modification, are permitted provided that the following conditions are met:
#
#   - Redistributions of source code must retain this notice, this list of 
#     conditions and the following disclaimer.
#
#   - Redistributions in binary form must provide access to this notice, this 
#     list of conditions and the following disclaimer, and the underlying source code.
#
#   - All modifications to this software must be clearly documented, and are 
#     solely the responsibility of the agent making the modifications.
#
#   - If significant modifications or enhancements are made to this software, 
#     the GSD Software Policy Manager (softwaremgr.fsl@noaa.gov) should be notified.
#
#   THIS SOFTWARE AND ITS DOCUMENTATION ARE IN THE PUBLIC DOMAIN AND ARE 
#   FURNISHED "AS IS." THE AUTHORS, THE UNITED STATES GOVERNMENT, ITS INSTRUMENTALITIES,
#   OFFICERS, EMPLOYEES, AND AGENTS MAKE NO WARRANTY, EXPRESS OR IMPLIED, AS TO
#   THE USEFULNESS OF THE SOFTWARE AND DOCUMENTATION FOR ANY PURPOSE. THEY ASSUME 
#   NO RESPONSIBILITY (1) FOR THE USE OF THE SOFTWARE AND DOCUMENTATION; OR (2) TO PROVIDE
#   TECHNICAL SUPPORT TO USERS. 
#############################################################################################
#############################################################################################

import os

try:
    import numpy as np
except:
    raise ImportError('The numpy library cannot be found!')

try:    
    from cdar.models.Shallow_Water_2D.ShallowWater2DModel import ShallowWaterModelClass as model
except:
    raise ImportError('The model library cannot be found!')

try:
    from cdar.utilities import namelist_read as nml
except:
    raise ImportError('The nml library cannot be found!')

try:
    import ESMF
except:
    raise ImportError('The ESMP library cannot be found!')
    
SampleTime = 1800

DataFilePath='./data'  # Path to store or read FG, B and OBS

# Read in parameters for model configurations
params = nml.namelist_read(DataFilePath)

# start up ESMF
# this call is not necessary unless you want to to override the
# default options:
#  LogKind = NONE
#  debug = False
manager = ESMF.Manager(logkind=ESMF.LogKind.SINGLE, debug=True)

# inquire for rank and proc from ESMF Virtual Machine
localPet = manager.local_pet
petCount = manager.pet_count
print "localPet = %d and petCount = %d " % (localPet, petCount)

# opening remarks
if localPet == 0:
    print "Welcome to the GEN true BE Utility!"
           
# Set up the model
domain = model(params)  

# Print domain info 
if localPet == 0:
    print domain  

# Read FG file
if localPet == 0:
    print "Reading First guess from file: %s" %  os.path.join(DataFilePath,'FG_'+domain.Name+'_'+str(SampleTime)+'.npz')
    npzfile = np.load(os.path.join(DataFilePath,'FG_'+domain.Name+'_'+str(SampleTime)+'.npz'))
    FG = npzfile['FG']
    FG = domain.strip_off_boundaries(FG)
    #u,v,h = np.split(FG,domain.VarDim)
    #vis.plot_3d(domain.x,domain.y,h,'FG')

if localPet == 0:
    print "Reading truth from file: %s" %  os.path.join(DataFilePath,'TRUTH_'+domain.Name+'_'+str(SampleTime)+'.npz')
    npzfile = np.load(os.path.join(DataFilePath,'TRUTH_'+domain.Name+'_'+str(SampleTime)+'.npz' ))
    Truth = npzfile['TRUTH']
    Truth = domain.strip_off_boundaries(Truth)
    #u,v,h = np.split(Truth,domain.VarDim)
    #vis.plot_3d(domain.x,domain.y,h,'Truth')

# Calculate the error
error = (FG - Truth).reshape(Truth.size)
  
# Calculate the B maxtrix
B = np.dot ( error[np.newaxis].transpose(), error[np.newaxis])

di = np.diag_indices(error.size)

B[di] = B[di] + 0.0000001
    
Binv = np.linalg.solve(B,np.eye(error.size))
        
print 'Is the Binv good ? ', np.allclose( np.dot(B, Binv), np.eye(int(np.sqrt(B.size))) )

if localPet == 0:
    print "The Binv size is %d x %d" % ((int(np.sqrt(B.size)),)*2)
    print "Save the BE in file: %s" %  os.path.join(DataFilePath,'B_'+domain.Name+'_'+'Truth'+'_'+str(domain.nx)+'_'+str(SampleTime))+'.npz'
    np.savez(os.path.join(DataFilePath,'B_'+domain.Name+'_'+'Truth'+'_'+str(domain.nx)+'_'+str(SampleTime)), B=Binv)

